import cv2
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image, ImageDraw, ImageFont
import math


class FengShuiAnalyzer:
    def __init__(self):
        # 风水基本规则定义
        self.rules = {
            "door_kitchen": {
                "description": "厨房门与入户门对冲",
                "is_bad": True,
                "remedy": "可在两门之间设置屏风或绿植"
            },
            "door_bathroom": {
                "description": "卫生间门与入户门对冲",
                "is_bad": True,
                "remedy": "可在两门之间设置屏风或安装门帘"
            },
            "kitchen_bedroom": {
                "description": "厨房与卧室相邻或对冲",
                "is_bad": True,
                "remedy": "可在厨房与卧室之间增加隔音隔热材料"
            },
            "bathroom_bedroom": {
                "description": "卫生间与卧室相邻或对冲",
                "is_bad": True,
                "remedy": "卫生间门常闭，保持干燥通风"
            },
            "bedroom_center": {
                "description": "卧室位于房屋中心",
                "is_bad": True,
                "remedy": "可考虑将卧室位置调整到房屋角落"
            },
            "straight_path": {
                "description": "入户门正对窗户形成穿堂煞",
                "is_bad": True,
                "remedy": "可在中间设置屏风或摆放家具阻挡"
            },
            "kitchen_center": {
                "description": "厨房位于房屋中心",
                "is_bad": True,
                "remedy": "厨房不宜设在房屋中心位置"
            }
        }

        # 房间类型映射
        self.room_types = {
            0: "unknown",
            1: "living_room",
            2: "bedroom",
            3: "kitchen",
            4: "bathroom",
            5: "door",
            6: "window",
            7: "balcony"
        }

        # 房间类型中文名称
        self.room_type_names = {
            "living_room": "客厅",
            "bedroom": "卧室",
            "kitchen": "厨房",
            "bathroom": "卫生间",
            "door": "门",
            "window": "窗",
            "balcony": "阳台",
            "unknown": "未知"
        }

    def load_image(self, image_path):
        """加载并预处理户型图"""
        img = cv2.imread(image_path)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        return img

    def detect_rooms(self, img):
        """简单的房间检测（实际应用中可能需要更复杂的计算机视觉算法）"""
        # 转为灰度图
        gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)

        # 二值化
        _, binary = cv2.threshold(gray, 200, 255, cv2.THRESH_BINARY_INV)

        # 形态学操作，去除噪点
        kernel = np.ones((3, 3), np.uint8)
        binary = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, kernel)

        # 查找轮廓
        contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

        # 筛选有效房间轮廓
        rooms = []
        for contour in contours:
            area = cv2.contourArea(contour)
            if area > 1000:  # 面积阈值，过滤噪点
                x, y, w, h = cv2.boundingRect(contour)
                rooms.append({
                    'type': 0,  # 默认未知类型
                    'x': x,
                    'y': y,
                    'width': w,
                    'height': h,
                    'center': (x + w / 2, y + h / 2),
                    'area': area
                })

        return rooms

    def classify_rooms(self, img, rooms):
        """房间类型分类（实际应用中可能需要机器学习模型）"""
        # 这里简化处理，实际应用中可以用颜色特征、形状特征或机器学习模型
        # 来更准确地识别房间类型
        for room in rooms:
            x, y, w, h = room['x'], room['y'], room['width'], room['height']
            roi = img[y:y + h, x:x + w]

            # 简单的颜色分析来判断房间类型
            # 实际应用中需要更复杂的算法
            avg_color = np.mean(roi, axis=(0, 1))

            # 这里使用简化的规则来模拟房间类型分类
            if avg_color[0] > 200 and avg_color[1] < 150 and avg_color[2] < 150:
                room['type'] = 3  # 厨房（红色区域）
            elif avg_color[0] < 150 and avg_color[1] > 200 and avg_color[2] < 150:
                room['type'] = 4  # 卫生间（绿色区域）
            elif avg_color[0] < 150 and avg_color[1] < 150 and avg_color[2] > 200:
                room['type'] = 2  # 卧室（蓝色区域）
            elif w > 1.5 * h or h > 1.5 * w:
                room['type'] = 5  # 门（长方形）
            else:
                room['type'] = 1  # 客厅（默认）

        return rooms

    def calculate_room_relationships(self, rooms):
        """计算房间之间的关系"""
        relationships = []

        # 检测房间之间的距离和位置关系
        for i, room1 in enumerate(rooms):
            for j, room2 in enumerate(rooms):
                if i != j:
                    # 计算两个房间中心点的距离
                    dx = room1['center'][0] - room2['center'][0]
                    dy = room1['center'][1] - room2['center'][1]
                    distance = math.sqrt(dx * dx + dy * dy)

                    # 检查是否对冲（两个房间中心点连线是否接近直线）
                    is_opposite = False
                    if distance > 0:
                        angle = abs(math.atan2(dy, dx) * 180 / math.pi)
                        # 水平或垂直对冲
                        is_opposite = (angle < 10 or angle > 170 or abs(angle - 90) < 10)

                    relationships.append({
                        'room1': i,
                        'room2': j,
                        'distance': distance,
                        'is_opposite': is_opposite
                    })

        return relationships

    def analyze_fengshui(self, rooms, relationships):
        """根据风水规则分析户型"""
        results = []

        # 检查门与厨房的关系
        doors = [i for i, room in enumerate(rooms) if room['type'] == 5]
        kitchens = [i for i, room in enumerate(rooms) if room['type'] == 3]

        for door in doors:
            for kitchen in kitchens:
                for rel in relationships:
                    if (rel['room1'] == door and rel['room2'] == kitchen) or \
                            (rel['room1'] == kitchen and rel['room2'] == door):
                        if rel['is_opposite'] and rel['distance'] < 100:  # 距离阈值
                            results.append({
                                'rule': 'door_kitchen',
                                'severity': 'high',
                                'description': self.rules['door_kitchen']['description'],
                                'remedy': self.rules['door_kitchen']['remedy']
                            })

        # 检查门与卫生间的关系
        bathrooms = [i for i, room in enumerate(rooms) if room['type'] == 4]

        for door in doors:
            for bathroom in bathrooms:
                for rel in relationships:
                    if (rel['room1'] == door and rel['room2'] == bathroom) or \
                            (rel['room1'] == bathroom and rel['room2'] == door):
                        if rel['is_opposite'] and rel['distance'] < 100:
                            results.append({
                                'rule': 'door_bathroom',
                                'severity': 'high',
                                'description': self.rules['door_bathroom']['description'],
                                'remedy': self.rules['door_bathroom']['remedy']
                            })

        # 检查厨房与卧室的关系
        bedrooms = [i for i, room in enumerate(rooms) if room['type'] == 2]

        for kitchen in kitchens:
            for bedroom in bedrooms:
                for rel in relationships:
                    if (rel['room1'] == kitchen and rel['room2'] == bedroom) or \
                            (rel['room1'] == bedroom and rel['room2'] == kitchen):
                        if rel['distance'] < 100:
                            results.append({
                                'rule': 'kitchen_bedroom',
                                'severity': 'medium',
                                'description': self.rules['kitchen_bedroom']['description'],
                                'remedy': self.rules['kitchen_bedroom']['remedy']
                            })

        # 检查卫生间与卧室的关系
        for bathroom in bathrooms:
            for bedroom in bedrooms:
                for rel in relationships:
                    if (rel['room1'] == bathroom and rel['room2'] == bedroom) or \
                            (rel['room1'] == bedroom and rel['room2'] == bathroom):
                        if rel['distance'] < 100:
                            results.append({
                                'rule': 'bathroom_bedroom',
                                'severity': 'medium',
                                'description': self.rules['bathroom_bedroom']['description'],
                                'remedy': self.rules['bathroom_bedroom']['remedy']
                            })

        # 检查房间是否位于中心
        image_center = (img.shape[1] / 2, img.shape[0] / 2)
        center_threshold = min(img.shape[1], img.shape[0]) * 0.2  # 中心区域阈值

        for i, room in enumerate(rooms):
            dx = room['center'][0] - image_center[0]
            dy = room['center'][1] - image_center[1]
            distance_to_center = math.sqrt(dx * dx + dy * dy)

            if distance_to_center < center_threshold:
                if room['type'] == 2:  # 卧室
                    results.append({
                        'rule': 'bedroom_center',
                        'severity': 'high',
                        'description': self.rules['bedroom_center']['description'],
                        'remedy': self.rules['bedroom_center']['remedy']
                    })
                elif room['type'] == 3:  # 厨房
                    results.append({
                        'rule': 'kitchen_center',
                        'severity': 'high',
                        'description': self.rules['kitchen_center']['description'],
                        'remedy': self.rules['kitchen_center']['remedy']
                    })

        # 检查穿堂煞（门对窗）
        windows = [i for i, room in enumerate(rooms) if room['type'] == 6]

        for door in doors:
            for window in windows:
                for rel in relationships:
                    if (rel['room1'] == door and rel['room2'] == window) or \
                            (rel['room1'] == window and rel['room2'] == door):
                        if rel['is_opposite'] and rel['distance'] > 200:  # 距离足够远
                            results.append({
                                'rule': 'straight_path',
                                'severity': 'high',
                                'description': self.rules['straight_path']['description'],
                                'remedy': self.rules['straight_path']['remedy']
                            })

        return results

    def visualize_results(self, img, rooms, results):
        """可视化分析结果"""
        # 创建一个副本用于绘制
        result_img = img.copy()

        # 绘制房间边界和类型
        for room in rooms:
            x, y, w, h = room['x'], room['y'], room['width'], room['height']
            room_type = self.room_types[room['type']]
            room_name = self.room_type_names[room_type]

            # 根据房间类型选择颜色
            if room_type == 'living_room':
                color = (255, 255, 0)  # 黄色
            elif room_type == 'bedroom':
                color = (0, 0, 255)  # 蓝色
            elif room_type == 'kitchen':
                color = (255, 0, 0)  # 红色
            elif room_type == 'bathroom':
                color = (0, 255, 0)  # 绿色
            elif room_type == 'door':
                color = (128, 128, 128)  # 灰色
            elif room_type == 'window':
                color = (128, 128, 255)  # 紫色
            else:
                color = (255, 255, 255)  # 白色

            # 绘制矩形和文本
            cv2.rectangle(result_img, (x, y), (x + w, y + h), color, 2)

            # 使用PIL添加中文文本
            pil_img = Image.fromarray(result_img)
            draw = ImageDraw.Draw(pil_img)

            try:
                # 尝试加载中文字体
                font = ImageFont.truetype("simhei.ttf", 14)
            except IOError:
                # 如果找不到中文字体，使用默认字体
                font = ImageFont.load_default()

            draw.text((x + 5, y + 5), room_name, font=font, fill=tuple(color))
            result_img = np.array(pil_img)

        # 绘制风水问题和建议
        y_offset = 20
        for result in results:
            severity = result['severity']
            description = result['description']
            remedy = result['remedy']

            # 根据严重程度选择颜色
            if severity == 'high':
                text_color = (255, 0, 0)  # 红色
            elif severity == 'medium':
                text_color = (255, 165, 0)  # 橙色
            else:
                text_color = (255, 255, 0)  # 黄色

            # 使用PIL添加文本
            pil_img = Image.fromarray(result_img)
            draw = ImageDraw.Draw(pil_img)

            try:
                # 尝试加载中文字体
                font = ImageFont.truetype("simhei.ttf", 12)
            except IOError:
                # 如果找不到中文字体，使用默认字体
                font = ImageFont.load_default()

            # 绘制问题描述
            draw.text((20, y_offset), f"问题: {description}", font=font, fill=text_color)
            y_offset += 20

            # 绘制建议
            draw.text((20, y_offset), f"建议: {remedy}", font=font, fill=(0, 255, 0))
            y_offset += 30

            result_img = np.array(pil_img)

        return result_img

    def analyze(self, image_path):
        """执行完整的风水分析流程"""
        # 加载图像
        img = self.load_image(image_path)

        # 检测房间
        rooms = self.detect_rooms(img)

        # 分类房间类型
        rooms = self.classify_rooms(img, rooms)

        # 计算房间关系
        relationships = self.calculate_room_relationships(rooms)

        # 分析风水
        results = self.analyze_fengshui(rooms, relationships)

        # 可视化结果
        result_img = self.visualize_results(img, rooms, results)

        return {
            'rooms': rooms,
            'results': results,
            'image': result_img
        }


# 使用示例
if __name__ == "__main__":
    analyzer = FengShuiAnalyzer()

    # 替换为你的户型图路径
    image_path = "house_layout.png"

    try:
        # 执行分析
        result = analyzer.analyze(image_path)

        # 显示结果
        plt.figure(figsize=(12, 8))
        plt.imshow(result['image'])
        plt.title("户型风水分析结果")
        plt.axis('off')
        plt.show()

        # 打印文字结果
        print("\n=== 风水分析结果 ===")
        for i, r in enumerate(result['results']):
            print(f"{i + 1}. {r['description']}")
            print(f"   严重程度: {r['severity']}")
            print(f"   建议: {r['remedy']}")
            print()

    except Exception as e:
        print(f"分析过程中发生错误: {e}")
